package com.tengju.data.domain.shared.tree;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.tengju.data.domain.shared.IdObject;
import lombok.Getter;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Getter
public class StableTree<T extends TreeData> extends Tree<T> {

    @JsonIgnore
    private final Map<Integer, Set<TreeNode<T>>> depthDictionary = new ConcurrentHashMap<>();

    private int maxDepth;

    public StableTree(TreeNode<T> root) {
        super(root);
        maxDepth = traverse(root, 0);
    }

    private int traverse(TreeNode<T> node, int currentDepth) {
        int depth = currentDepth;
        //更新节点深度
        node.updateDepth(depth);
        //更新深度字典
        this.allTreeNodes.put(node.getId(), node);
        getDepthDictionary().computeIfAbsent(node.getDepth(), k -> new HashSet<>());
        Set<TreeNode<T>> treeNodeSet = getDepthDictionary().get(node.getDepth());
        if (CollectionUtils.isEmpty(treeNodeSet)){
            treeNodeSet = new HashSet<>();
            treeNodeSet.add(node);
            getDepthDictionary().put(node.getDepth(), treeNodeSet);
        } else {
            treeNodeSet.add(node);
        }
        //更新全部节点字典
        this.allTreeNodes.put(node.getId(), node);
        if (node.getChildren() != null) {
            depth = currentDepth + 1;
            for (TreeNode<T> child : node.getChildren()) {
                int childDepth = traverse(child, currentDepth + 1);
                if (childDepth > depth) {
                    depth = childDepth;
                }
            }
        }
        return depth;
    }


    public Set<TreeNode<T>> findNodes(int depth) {
        return Optional.ofNullable(depthDictionary.get(depth))
                .map(Collections::unmodifiableSet)
                .orElse(null);
    }


    public int countNodes() {
        return getAllTreeNodes().size();
    }

//    @Override
//    void addNodes(TreeNode<T> parent, Set<TreeNode<T>> children) {
//        super.addNodes(parent, children);
//        synchronized (getRoot()) {
//            //更新节点字典
//            recordNodes(children);
//            //更新树深度
//            if (parent.maxDepth() > maxDepth) {
//                maxDepth = parent.maxDepth();
//            }
//        }
//    }
//
//    private void recordNodes(Set<TreeNode<T>> nodes) {
//        nodes.forEach(node -> {
//            updateDictionary(node);
//            if (node.getChildren() != null) {
//                recordNodes(node.getChildren());
//            }
//        });
//    }



    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        if (!super.equals(o)) return false;
        StableTree<?> that = (StableTree<?>) o;
        return getRoot() == that.getRoot();
    }

    @Override
    public int hashCode() {
        return Objects.hash(super.hashCode(), getRoot());
    }
}
